/*!
  @file mapping.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include <ros/ros.h>
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>
#include <tf/transform_listener.h>
#include <std_msgs/Float32.h>
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include "Labeling/Labeling.h"

namespace mrover 
{

/*! 
 *  @brief  An awesome class
 *  @author Kyohei Otsu <kyon@ac.jaxa.jp>
 *  @date   2013
 *  
 *  This is a detailed description of the class 
 */

class Mapping 
{
 public:
  Mapping(ros::NodeHandle &nh) 
      : nh_(nh)
      , private_nh_("~")
      , it_(nh_)
      , scale_(0.0)
      , offset_(0)
      , sigma_(3.0)
      , speckle_size_(10)
  {
    // initializing pub/sub, server
    img_sub_ = it_.subscribeCamera(nh_.resolveName("image_in"), 3, &Mapping::imageCB, this);
    img_pub_ = it_.advertise(nh_.resolveName("image_out"), 3);
    sigma_sub_ = nh_.subscribe(nh_.resolveName("sigma_in"), 3, &Mapping::sigmaCB, this);
    speckle_sub_ = nh_.subscribe(nh_.resolveName("speckle_in"), 3, &Mapping::speckleCB, this);
    marker_pub_ = nh_.advertise<visualization_msgs::MarkerArray>(nh_.resolveName("obstacle_out"), 3);

    // initializing optional parameters
    // e.g.) private_nh_.param<TYPE>("PARAM_NAME", PARAM_VAR, DEFAULT_VALUE);
  }

  ~Mapping() { }

  void imageCB(const sensor_msgs::ImageConstPtr &img_msg, const sensor_msgs::CameraInfoConstPtr &cam_info_msg);

  void sigmaCB(const std_msgs::Float32ConstPtr &msg) { sigma_ = msg->data; }
  void speckleCB(const std_msgs::Float32ConstPtr &msg) { speckle_size_ = msg->data; }


 private:

  class Obstacle
  {
   public:
    Obstacle() : id_(0) { }
    ~Obstacle() { }

    visualization_msgs::Marker &getMarker()
    {
      if (!marker_) generateMarker();
      return *marker_;
    }
    visualization_msgs::Marker &deleteMarker()
    {
      if (!marker_) generateMarker();
      marker_->action = visualization_msgs::Marker::DELETE;
      return *marker_;
    }

    void setID(int id) { id_ = id; }
    void setXY(int x, int y)
    {
      position_.x = x;
      position_.y = y;
    }
   private:
    void generateMarker()
    {
      marker_.reset(new visualization_msgs::Marker);
      marker_->header.stamp = ros::Time::now();
      marker_->header.frame_id = "/rover/base_link";
      marker_->id = id_;
      marker_->type = visualization_msgs::Marker::CYLINDER;
      marker_->action = visualization_msgs::Marker::ADD;
      marker_->pose.position.x = position_.x;
      marker_->pose.position.y = position_.y;
      marker_->pose.position.z = 0;
      marker_->pose.orientation.x = 0.0;
      marker_->pose.orientation.y = 0.0;
      marker_->pose.orientation.z = 0.0;
      marker_->pose.orientation.w = 1.0;
      marker_->scale.x = 0.2;
      marker_->scale.y = 0.2;
      marker_->scale.z = 0.4;
      marker_->color.a = 1.0;
      marker_->color.r = 1.0;
      marker_->color.g = 0.0;
      marker_->color.b = 0.0;
    }
    visualization_msgs::MarkerPtr marker_;
    size_t id_;
    cv::Point2f position_;
  };


  void process();
  void obtainHomographyMatrix(const cv::Mat &K);
  
  ros::NodeHandle nh_;
  ros::NodeHandle private_nh_;
  image_transport::ImageTransport it_;
  image_transport::CameraSubscriber img_sub_;
  image_transport::Publisher img_pub_;
  ros::Publisher marker_pub_;
  ros::Subscriber sigma_sub_;
  ros::Subscriber speckle_sub_;

  LabelingSS lss;
  int speckle_size_;

  cv::Mat image_;
  float scale_;
  int offset_;
  double sigma_;

  cv::Mat H_;
};

void Mapping::imageCB(const sensor_msgs::ImageConstPtr &img_msg, const sensor_msgs::CameraInfoConstPtr &cam_info_msg)
{
  cv_bridge::CvImage::Ptr cv_ptr;
  try {
    cv_ptr = cv_bridge::toCvCopy(img_msg, "mono8");
  }
  catch(cv_bridge::Exception &e) {
    ROS_ERROR("cv_bridge: conversion error");
    return;
  }

  /* resize image */
  if (!scale_)
  {
    switch(cv_ptr->image.cols)
    {
      case 320:   scale_ = 0.5;   break;
      case 640:   scale_ = 0.25;  break;
      case 1024:  scale_ = 0.25;  break;
      default:    scale_ = 0.25;  break;
    }
    offset_ = 0.25 * scale_ * cv_ptr->image.rows;
  }
  int w = scale_ * cv_ptr->image.cols;
  int h = scale_ * cv_ptr->image.rows;
  if (!image_.data) image_.create(h, w, CV_8U);
  cv::resize(cv_ptr->image, image_, image_.size());

  cv::Mat K(3, 3, CV_32F);
  for (int i = 0; i < 3; ++i)
    for (int j = 0; j < 3; ++j)
      K.at<float>(i, j) = cam_info_msg->K[i*3+j];
  obtainHomographyMatrix(K);

  /* processing */
  process();
}

void Mapping::obtainHomographyMatrix(const cv::Mat &K)
{
  if (!H_.data) H_.create(3, 3, CV_32F);
  static tf::TransformListener listener;
  tf::StampedTransform transform;
  try
  {
    listener.waitForTransform("/rover/left_camera", "/rover/base_link", ros::Time::now(), ros::Duration(3.0));
    listener.lookupTransform("/rover/left_camera", "/rover/base_link", ros::Time(0), transform);
  }
  catch (tf::TransformException &e)
  {
    ROS_WARN("tf not published yet. retrying...");
    return;
  }
  tf::Vector3    position = transform.getOrigin();
  tf::Quaternion quaternion = transform.getRotation();
  tf::Matrix3x3  orientation(quaternion);

  cv::Mat RT(3, 4, CV_32F);
  for (int i = 0; i < 3; ++i)
    for (int j = 0; j < 3; ++j)
      RT.at<float>(i, j) = orientation[i][j];
  RT.at<float>(0, 3) = position.x();
  RT.at<float>(1, 3) = position.y();
  RT.at<float>(2, 3) = position.z();

  cv::Mat P = K * RT;
  cv::Mat H0 = H_.col(0);
  cv::Mat H1 = H_.col(1);
  cv::Mat H2 = H_.col(2);
  P.col(0).copyTo(H0);
  P.col(1).copyTo(H1);
  P.col(3).copyTo(H2);
  H_ /= H_.at<float>(2, 2);  // normalize
  H_ = H_.inv();
}

void Mapping::process()
{
  cv::Mat mask = cv::Mat::ones(image_.size(), CV_8U);
  cv::rectangle(mask, cv::Point(0, 0), cv::Point(image_.cols, offset_-1), 0, -1);
  cv::Mat mean(1, 1, CV_64F);
  cv::Mat stddev(1, 1, CV_64F);
  cv::meanStdDev(image_, mean, stddev, mask);
  
  float m = mean.at<double>(0, 0);
  float s = stddev.at<double>(0, 0);

  cv::Mat bin_image(image_.size(), CV_8U);
  cv::threshold(image_, bin_image, m + sigma_ * s, 0, cv::THRESH_TOZERO_INV);
  cv::threshold(bin_image, bin_image, m - sigma_ * s, 255, cv::THRESH_BINARY);
  bin_image = ~bin_image;

  // debug
  {
    cv_bridge::CvImagePtr cv_ptr(new cv_bridge::CvImage);
    cv_ptr->image = bin_image;
    cv_ptr->header.stamp = ros::Time::now();
    cv_ptr->header.frame_id = "/mrover/left_camera";
    cv_ptr->encoding = "mono8";
    img_pub_.publish(cv_ptr->toImageMsg());
  }

  /* labeling */
  short buf[image_.cols * (image_.rows - offset_)];
  for (int i = 0; i < image_.rows - offset_; ++i)
    for (int j = 0; j < image_.cols; ++j)
      buf[i * image_.cols + j] = image_.at<uchar>(i + offset_, j);
  lss.Exec(buf, buf, image_.cols, image_.rows - offset_, true, speckle_size_);

  ROS_INFO("Num obstacles: %d (%d pixel speckles eliminated)", lss.GetNumOfResultRegions(), speckle_size_);

  RegionInfoSS *ri;
  cv::Point min_pt, max_pt;
  std::vector<cv::Point2f> image_pts;
  std::vector<cv::Point2f> obj_pts;
  for (int i = 0; i < lss.GetNumOfResultRegions(); ++i)
  {
    ri = lss.GetResultRegionInfo(i);
    ri->GetMin(min_pt.x, min_pt.y);
    ri->GetMax(max_pt.x, max_pt.y);
    cv::Point2f lpt(min_pt.x / scale_, (max_pt.y + offset_) / scale_);
    cv::Point2f rpt(max_pt.x / scale_, (max_pt.y + offset_) / scale_);
    cv::Point2f cpt(0.5 * (lpt.x + rpt.x), 0.5 * (lpt.y + rpt.y));
    //image_pts.push_back(lpt);
    //image_pts.push_back(rpt);
    image_pts.push_back(cpt);
    //ROS_INFO_STREAM(cpt);
  }

  cv::Mat disp = image_.clone();
  for (int i = 0; i < image_pts.size(); ++i)
  {
    cv::Point pt(image_pts[i].x*scale_, image_pts[i].y*scale_);
    cv::circle(disp, pt, 30, 0, 2);
  }
  //cv::imshow("dsip", disp);
  cv::waitKey(1);


  if (image_pts.size() > 0) cv::perspectiveTransform(image_pts, obj_pts, H_);
  //std::vector<Obstacle> obstacles(obj_pts.size());
  const int MAX_OBJ = 30;
  std::vector<Obstacle> obstacles(30);
  visualization_msgs::MarkerArrayPtr marker_a(new visualization_msgs::MarkerArray);
#if 0
  for (int i = 0; i < obj_pts.size(); ++i)
  {
    obstacles[i].setID(i);
    obstacles[i].setXY(obj_pts[i].x, obj_pts[i].y);
    ROS_INFO_STREAM(obj_pts[i]);
    marker_a->markers.push_back(obstacles[i].getMarker());
  }
#endif
  for (int i = 0; i < MAX_OBJ; ++i)
  {
    obstacles[i].setID(i);
    if (i < obj_pts.size()) 
    {
      //ROS_INFO_STREAM(obj_pts[i]);
      obstacles[i].setXY(obj_pts[i].x, obj_pts[i].y);
      marker_a->markers.push_back(obstacles[i].getMarker());
    }
    else
    {
      marker_a->markers.push_back(obstacles[i].deleteMarker());
    }
  }
  marker_pub_.publish(marker_a);
}

}  // namespace mrover

int main(int ac, char **av)
{
  ros::init(ac, av, "Mapping");
  ros::NodeHandle nh;
  mrover::Mapping handler(nh);
  ros::spin();
  return 0;
}

